#include <QCoreApplication>
#include <QDebug>
#include "SM4.h"
#include "SM3.h"
#include "SM2.h"

void createTestData(QByteArray& data, int size)
{
    data.resize(size);
    for (int i = 0; i < size; i++)
    {
        data[i] = i % 128;
    }
}

void testSM4(const QByteArray& data)
{
    qDebug() << "------------------------SM4 Test------------------------";
    QByteArray plainText = data; // 明文
    QByteArray encryptText;      // 密文
    QByteArray decryptText;      // 解密内容

    QByteArray key = QByteArray::fromHex("8cc72b05705d5c46f412af8cbed55aad");  // 秘钥
    QByteArray ivec = QByteArray::fromHex("667b02a85c61c786def4521b060265e8"); // 初始化向量

    // ecb模式加解密验证
    SM4 sm4;
    sm4.ecb_encrypt(plainText, encryptText, key, true);    // 加密
    sm4.ecb_encrypt(encryptText, decryptText, key, false); // 解密
    qDebug() << "SM4 ecb encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // cbc模式加解密验证
    sm4.cbc_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.cbc_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 cbc encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // cfb1模式加解密验证
    sm4.cfb1_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.cfb1_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 cfb1 encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // cfb8模式加解密验证
    sm4.cfb8_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.cfb8_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 cfb8 encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // cfb128模式加解密验证
    sm4.cfb128_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.cfb128_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 cfb128 encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // ofb模式加解密验证
    sm4.ofb_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.ofb_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 ofb encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // ctr模式加解密验证
    sm4.ctr_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.ctr_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 ctr encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    // gcm模式加解密验证
    sm4.gcm_encrypt(plainText, encryptText, key, ivec, true);    // 加密
    sm4.gcm_encrypt(encryptText, decryptText, key, ivec, false); // 解密
    qDebug() << "SM4 gcm encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");
    encryptText.clear();
    decryptText.clear();

    qDebug() << "";
}

void testSM3(const QByteArray& data)
{
    qDebug() << "------------------------SM3 Test------------------------";
    QByteArray plainText = data; // 明文
    QByteArray digestText;       // 摘要

    SM3 sm3;
    digestText = sm3.digest(plainText); // 计算摘要
    qDebug() << "SM3 digest:" << digestText.toHex();

    qDebug() << "";
}

void testSM2(const QByteArray& data)
{
    qDebug() << "------------------------SM2 Test------------------------";

    // 测试生成秘钥对
    QByteArray priKey, pubKey;
    SM2 sm2;
    sm2.generateKeyPair(priKey, pubKey);
    qDebug() << "private key:" << priKey;
    qDebug() << "public key:" << pubKey;

    // 测试生成秘钥对文件
    QString curDir = qApp->applicationDirPath();
    bool ret = sm2.generateKeyPair(curDir + "/pri.pem", curDir + "/pub.pem");
    qDebug() << "SM2 key pair file generate" << (ret ? "succeeded" : "failed");

    // 测试对数据进行加解密
    QByteArray plainText = data;
    QByteArray encryptText;
    QByteArray decryptText;

    sm2.encrypt(plainText, encryptText, pubKey);   // 加密
    sm2.decrypt(encryptText, decryptText, priKey); // 解密
    qDebug() << "SM2 encrypt verify" << ((decryptText == plainText) ? "succeeded" : "failed");

    // 测试签名与验签
    QByteArray digest;
    QByteArray sign;
    SM3 sm3;
    digest = sm3.digest(plainText); // 计算摘要
    sm2.sign(digest, sign, priKey);         // 签名
    ret = sm2.verify(digest, sign, pubKey); // 验签
    qDebug() << "SM2 sign verify" << (ret ? "succeeded" : "failed");
}

int main(int argc, char *argv[])
{
    QCoreApplication a(argc, argv);

    // 产生64K-1的测试数据
    QByteArray data;
    createTestData(data, 65535);

    // 进行各算法测试
    testSM4(data);
    testSM3(data);
    testSM2(data);

    return a.exec();
}
